#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/7/27 1:23 上午
# @Author  : WangZhixing
import torch
from torch_geometric.nn import GCNConv
import torch.nn.functional as F

# 包含两个模型，
# 1、纯GCN2层模型
# 参数：
#     1：num_features输入层
#     2：hidden_channels中间层
#     3：num_classes输出层

class GCN(torch.nn.Module):
    def __init__(self, num_features, hidden_channels, num_classes):
        super(GCN, self).__init__()
        torch.manual_seed(12345)
        # 降维一下尝试nn.line

        # self.prec = torch.nn.Linear()

        self.conv1 = GCNConv(num_features, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, num_classes)

    def forward(self, x, edge_index, edge_attr):
        x = self.conv1(x, edge_index, edge_weight=edge_attr)
        x = x.relu()

        # 激活函数多改一些
        #####################################################
        x = F.dropout(x, p=0.5, training=self.training)
        #####################################################
        x = self.conv2(x, edge_index)
        return x

# 1、带线性层的降维GCN2层模型
# 参数：
#     1：num_features输入层
#     2：hidden_channels中间层
#     3：num_classes输出层


class Linear_GCN(torch.nn.Module):
    def __init__(self, num_features, hidden_channels1,hidden_channels2, num_classes):
        super(Linear_GCN, self).__init__()
        torch.manual_seed(12345)
        # 降维一下尝试nn.line
        self.pre = torch.nn.Linear(num_features,hidden_channels1)
        self.conv1 = GCNConv(hidden_channels1, hidden_channels2)
        self.conv2 = GCNConv(hidden_channels2, num_classes)


    def forward(self, x, edge_index, edge_attr):
        x = self.pre(x)
        x = self.conv1(x, edge_index, edge_weight=edge_attr)
        x = x.relu()

        # 激活函数多改一些

        #####################################################
        x = F.dropout(x, p=0.5, training=self.training)
        #####################################################
        x = self.conv2(x, edge_index)
        return x


